const builtin = @import("builtin");
const csr = @import("riscv").csr;
const root = @import("kernel.zig");
const clint = @import("clint.zig");
const lock = @import("lock.zig");

const print = root.print;

extern fn switch_to(to: *Context) callconv(.c) void; // 返回类型不能是 noreturn

const Context = packed struct {
    // Layout must be kept in sync with switch_to
    // ignore x0
    ra: usize = 0,
    sp: usize,
    gp: usize = 0,
    tp: usize = 0,
    t0: usize = 0,
    t1: usize = 0,
    t2: usize = 0,
    s0: usize = 0,
    s1: usize = 0,
    a0: usize = 0,
    a1: usize = 0,
    a2: usize = 0,
    a3: usize = 0,
    a4: usize = 0,
    a5: usize = 0,
    a6: usize = 0,
    a7: usize = 0,
    s2: usize = 0,
    s3: usize = 0,
    s4: usize = 0,
    s5: usize = 0,
    s6: usize = 0,
    s7: usize = 0,
    s8: usize = 0,
    s9: usize = 0,
    s10: usize = 0,
    s11: usize = 0,
    t3: usize = 0,
    t4: usize = 0,
    t5: usize = 0,
    t6: usize = 0,
    // upon is trap frame

    // save the pc to run in next schedule cycle
    pc: usize, // offset: 31 * @sizeOf(usize)

    pub fn new(thread_stack: []u8, func: fn () void) Context {
        return .{
            .sp = @intFromPtr(&thread_stack[thread_stack.len - 1]),
            .pc = @intFromPtr(&func),
        };
    }
};

const Task = struct {
    // 栈太小的话，会覆盖 context 的空间，mepc 地址不对导致执行 mret 异常
    const STACK_SIZE = if (builtin.mode == .Debug) 2048 else 512;

    context: Context,
    stack: [STACK_SIZE]u8 = @splat(65), // 这里的默认初始化也没有用到，two_tasks 是 undefined
};

// 为了让 switch_to 能访问到，它必须是全局变量
const MAX_TASKS = 2;
pub var two_tasks: [MAX_TASKS]Task = undefined; // 在堆里分配会好写一点吗？

var top: usize = 0;
var current: usize = 1;
var spinlock: lock.SpinLock = undefined;

pub fn info() void {
    for (&two_tasks, 0..) |*task, i| {
        print("task{d}: stack: {*} -> {*} {c}\n", .{
            i,
            &task.stack[0],
            &task.stack[Task.STACK_SIZE - 1],
            task.stack[0],
        });
    }
}

pub fn schedule() void { // 返回类型不能是 noreturn
    current = (current + 1) % top;
    switch_to(&two_tasks[current].context);
}

pub inline fn yield() void {
    // trigger a machine-level software interrupt
    const hart = 0;
    clint.software.trigger(hart);
}

pub fn init() void {
    // 初始化 mscratch
    csr.mscratch.write(0);

    clint.software.init();

    spinlock = .init();

    inline for (.{ user_task0, user_task1 }) |func| {
        two_tasks[top].context = .new(&two_tasks[top].stack, func);
        top += 1;
    }
}

fn user_task0() void {
    print("Task 0: Created!\n", .{});

    while (true) {
        var held = spinlock.acquire();

        print("Task 0: Begin ... \n", .{});

        for (0..5) |_| {
            print("Task 0: Running...\n", .{});
            delay(1000);
        }

        print("Task 0: End ... \n", .{});

        held.release();

        // 让另一 task 有机会执行
        delay(5000);
    }
}

fn user_task1() void {
    print("Task 1: Created!\n", .{});

    while (true) {
        var held = spinlock.acquire();

        print("Task 1: Begin ... \n", .{});
        for (0..5) |_| {
            print("Task 1: Running...\n", .{});
            delay(1000);
        }
        print("Task 1: End ... \n", .{});

        held.release();

        // 让另一 task 有机会执行
        delay(5000);
    }
}

fn delay(x: usize) void {
    var i: usize = 50000 * x;
    const ptr: *volatile usize = &i;
    while (ptr.* > 0) {
        ptr.* -= 1;
    }
}
